"""Low-level Maths. All functions here work on vectors, not timedata."""

import numpy as np

def CommonInit(y, size):
    """Common initialization used in a lot of the functions below."""
    n = len(y)
    newN = max(n-size+1, 0)
    if newN > 0:
        z = np.zeros(newN, dtype=float)
    else:
        z = np.array([], dtype=float)
    return (n, newN, z)
    

def EMA(y, size):
    """Exponential Moving Average.
    
    The result has length=n-size+1, where n is the size of the original vector.
    First EMA(0) value corresponds the time position of (size-1)
    
    EMA(0) = simple moving average of size values
    k = 2/(size+1)
    EMA(i) = k*y(i+size-1) + (1-k)*EMA(i-1)  (0 < i < n-size+1)
    
    References:
        https://en.wikipedia.org/wiki/Moving_average#Exponential_moving_average
    """
    (n, newN, z) = CommonInit(y, size)
    if newN > 0:
        # First calculated value is simple mean
        z[0] = np.mean(y[0:size])
        # Then starts calculating the EMA
        k = 2./(size+1)
        for i in range(1, newN):
            z[i] = (1-k)*z[i-1]+k*y[i+size-1]
    return z
        
def MaxN(y, size):
    """Maximum of last size points.
    """
    (n, newN, z) = CommonInit(y, size)
    if newN > 0:
        for i in range(0, newN):
            z[i] = np.max(y[i:i+size])
    return z
    
def MinN(y, size):
    """Maximum of last size points.
    """
    (n, newN, z) = CommonInit(y, size)
    if newN > 0:
        for i in range(0, newN):
            z[i] = np.min(y[i:i+size])
    return z

def Stochastic(close, high, low, size):
    """Stochastic oscillator.
    
    Arguments:
        y -- data vector
        size -- Formula parameter for MinN() and MaxN(), see formula
    
    According to [1]
    %K = 100[(Close - L14)/(H14 - L14)] 
    In this case, "14" is a typical value for the size parameter.
    
    References
    [1] http://www.investopedia.com/terms/s/stochasticoscillator.asp
    """
    
    (n, newN, z) = CommonInit(close, size)
    if newN > 0:
        mins = MinN(low, size)
        maxs = MaxN(high, size)
        z = (close[size-1:]-mins)/(maxs-mins)
    return z

def DesignMatrix(x, order):
    """Design matrix: first column of ones plus order columns of shifted signal.
    """
    numRows = len(x)-order+1
    if numRows < 1:
        raise RuntimeError("x (%d points) too short for %d columns!" % (len(x), order))
    r = np.zeros((numRows, order+1))
    r[:, 0] = 1 # Constant term
    for i in range(0, numRows):
#        print x[i:i+order-1]
        r[i,1:] = x[i:i+order]
    return r





def LSPredictions(signal, trainingSize, order, travel):
    """Calculates whole signal of predicted values.
    
    Arguments:
        signal -- signal
        trainingSize -- number of points to use from signal for each prediction.
                    if 0, will start from signal[0] at all times,
                    otherwise the 
        
        order -- number of columns of data matrix
        travel -- prediction time travel. 0: today; 1: tomorrow; -1: yesterday.
                  Never uses  INTERRUPTED
                  
    Output:
        z -- lagged signal. The amount of lag is trainingSize-1+travel
    """
    
    N = len(signal)
    nit_total = N-trainingSize+1
    if nit_total < 1:
        raise RuntimeError("not enough data: len(signal)=%d; trainingSize=%d" % (N, trainingSize))
    nit_train = nit_total-travel # Number of training iterations
    z = np.zeros(nit_total)
    if nit_train < 1:
        raise RuntimeError("not enought data: len(signal)=%d; order=%d; travel=%d" % (N, order, travel))
    for it in range(nit_total):
        if it < nit_train:
            (X, Y) = GetXY(signal, trainingSize, order, travel, it)
            coeff = np.linalg.lstsq(X, Y)[0]
        print "-0-", it
        print "-1-", GetInputSignal(signal, trainingSize, order, it)
        print "-2-", coeff
        z[it] = np.dot(GetInputSignal(signal, trainingSize, order, it), coeff)
    return z

def GetXY(signal, trainingSize, order, travel, i):
    """Design matrix (X) and output matrix (Y)."""
    X = DesignMatrix(signal[i:i+trainingSize], order)
    Y = signal[i+travel+order-1:i+trainingSize+travel]
    return (X, Y)

def GetInputSignal(signal, trainingSize, order, i, flagOffset=1):
    """Part of signal to be used as input signal for some prediction.
    
    Arguments:
        flagOffset=1 -- whether to include [1] as first element of output."""
    i0 = i+trainingSize-order
    x = signal[i0:i0+order]
    if flagOffset:
        return np.concatenate(([1], x))
    else:
        return x



def Momentum(signal, size=1):
    """"Momentum" indicator.
    
    Result is lagged by size.
    
    Arguments:
        size=1 -- size parameter. Size=1 results in simple difference between
                  current and previous value.
    
    References:
        [1] Jake Bernstein, The compleat day trader vol 1
    """
    n = len(signal)
    if size > n: return signal[:0] ###, "Size (%d) must be smaller than len(signal) (%d)" % (size, n)
    r = signal[size:]-signal[0:-size]
    return r
